#!/usr/bin/env python
#
# Copyright 2013 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Module for running various network-related tests on a testing device.

Functions in this script are used to run various network performance related
tests.
"""
__author__ = "Joey Freeland (joeyfreeland@gmail.com)"

import logging
import re
import time
import urllib
import urllib2

import device

SITE_CHECK = {"Bing": "2013 Microsoft", "Google": "Lucky", "Yahoo":
              "Yahoo! Indonesia"}


def CaptivePortalLogin(networks, ssid):
  """Login to the captive portal of the network being tested.

  Args:
    networks: A dictionary of credentials for the networks we're testing.
    ssid: The SSID of the network we want to connect to now.

  Returns:
    None
  """
  logging.info("Attempting to authenticate to %s.", ssid)
  if "sid" in networks[ssid].keys():
    load_portal = urllib2.Request("http://icanhas.cheezburger.com/")
    portal_output = urllib2.urlopen(load_portal, None, 15).read()
    sid_re = re.search(r"sid=(.*)\"", portal_output)
    networks[ssid]["sid"] = sid_re.group(1)
  if "FreeNet" in networks[ssid].keys():
    load_portal = urllib2.Request("http://icanhas.cheezburger.com/")
    portal_output = urllib2.urlopen(load_portal, None, 15).read()
    key_re = re.search(r"zairvai_freenet.*value=\"(.*)\"", portal_output)
    networks[ssid]["data"]["zairvai_freenet"] = key_re.group(1)
  if "INDOSATNET" in networks[ssid].keys():
    load_portal = urllib2.Request("http://icanhas.cheezburger.com/")
    portal_output = urllib2.urlopen(load_portal, None, 15).read()
    key_re = re.search(r"/pas/start?key=(.*)\"", portal_output)
    networks[ssid]["data"]["key"] = key_re.group(1)
  url = networks[ssid]["url"]
  data = urllib.urlencode(networks[ssid]["data"])
  if "post" in networks[ssid].keys():
    logging.info("Logging into captive portal via HTTP POST.")
    urllib2.urlopen(url, data, 15)
    return
  else:
    logging.info("Logging into captive portal via HTTP GET.")
    url = "%s?%s" % (url, data)
    urllib2.urlopen(url, None, 15)
    return


def CheckDNS(host, nameserver=""):
  """Resolve a hostname against a nameserver and record the latency.

  Args:
    host: FQDN that we want to resolve.
    nameserver: IP address of the nameserver to attempt to use for resolution.
                Can be none to default to system configured nameservers.

  Returns:
    Either a number that represents the amount of time (in ms) it took to
    resolve the name, or a 0 if we were unable to resolve the name.
  """
  if nameserver:
    logging.info("Checking DNS resolution using %s.", nameserver)
    dns_cmd = "/usr/bin/nslookup %s %s" % (host, nameserver)
  else:
    logging.info("Checking DNS resolution using DHCP assigned nameserver.")
    dns_cmd = "/usr/bin/nslookup %s" % host
  start = int(round(time.time() * 1000))
  results, errors = device.RunCommand(dns_cmd, errors=True)
  finish = int(round(time.time() * 1000))
  latency = finish - start
  if errors:
    logging.debug("Encountered problems resolving %s:\n\n%s\n\n", host, errors)
    return 0
  dns_search = re.search(r"(Name\:(.*\n){1,2}(Address.*([0-9]{1,3}\.){,3}"
                         "[0-9]{1,3}))", results)
  if dns_search:
    logging.info("Successfully resolved domain name in %d ms.", latency)
    return latency


def CheckRedirect():
  """Load a random website and check HTTP headers to see if we were redirected.

  This function is intended to be used when first connecting to a network.  This
  is intended to determine whether or not the network has a captive portal by
  looking for an HTTP redirect error code.

  Args:
    None

  Returns:
    Pass (1 = redirected as expected) or fail (0 = not redirected).
  """
  logging.info("Checking for captive portal.  Expecting HTTP redirect.")
  request = urllib2.Request("http://app-website/check_redirect")
  try:
    redirect = urllib2.urlopen(request, None, 15)
  except urllib2.HTTPError as e:
    logging.info("HTTP error code encountered: %s", e.code)
    return True
  response = redirect.read()
  redirect_check = re.search(r"You were not redirected.", response)
  if redirect_check:
    logging.info("No captive portal detected.")
    return
  else:
    logging.info("Captive portal detected.")
    return True


def LoadSite(site, name):
  """Load a website and verify that a specific term exists in the results.

  Args:
    site: Address of the website to be loaded.  This should probably be url.
    name: Nickname for the website being loaded.

  Returns:
    Returns the latency in ms that it took to load the page or 0 for failed
    attempt.
  """
  logging.info("Attempting to load %s", site)
  request = urllib2.Request(site)
  start = int(round(time.time() * 1000))
  response = urllib2.urlopen(request, None, 15)
  finish = int(round(time.time() * 1000))
  latency = finish - start
  html = response.read()
  site_audit = re.search(SITE_CHECK[name], html)
  if site_audit:
    logging.info("Successfully loaded %s in %d ms", name, latency)
    return latency
  else:
    logging.debug("Loaded an unexpected page while trying to load %s.", name)
    return 0


def GetExternalIP():
  """Load a webpage to determine the external IP address of the device."""
  request = urllib2.Request("http://app-website/getip")
  ip = urllib2.urlopen(request, None, 15).read()
  logging.info("External IP address is: %s", ip)
  return ip


def ReverseSshTunnel(ssh_host, ssh_port, ssh_user):
  """Create a reverse SSH tunnel to a bastion host."""
  # TODO(joeyfreeland): Need to add some way to check to see if a tunnel is in
  # use.
  logging.info("Creating reverse SSH tunnel.")
  cmd = "/usr/bin/ssh -f -N -R %s:localhost:22 %s@%s" % (ssh_port, ssh_user,
                                                         ssh_host)
  device.RunCommand(cmd, errors=False, results=False)


def RunNDT():
  """Run the Internet2 / M-Lab Network Diagnostic Tool client.

  Args:
    None

  Returns:
    ndt_up: A float representing the uplink speed in Mbps.
    ndt_down: A float representing the download speed in Mbps.
    # ip_address: The IP address that NDT recognized.
  """
  logging.info("Running the NDT web100 client.")
  cmd = "/usr/local/bin/web100clt -n ndt.iupui.donar.measurement-lab.org"
  results, errors = device.RunCommand(cmd, timeout=45, errors=True)
  if errors:
    logging.debug("Encountered a system error while running NDT.")
    return 0, 0
  up_search = re.search(r"\(client to server\).* (\d+\.\d+) (\w\w)\/s", results)
  down_search = re.search(r"\(server to client\).* (\d+\.\d+) (\w\w)\/s",
                          results)
  if up_search and down_search:
    up = up_search.group(1)
    up_unit = up_search.group(2)

    down = down_search.group(1)
    down_unit = down_search.group(2)
    logging.info("NDT results: %s %s up and %s %s down", up, up_unit, down,
                 down_unit)
    if up_unit == "Mb": up = float(up) * 1048576
    else: up = float(up) * 1024
    if down_unit == "Mb": down = float(down) * 1048576
    else: down = float(down) * 1024
    return int(up), int(down)
  else:
    logging.debug("Encountered a problem while running the web100 client.")
    logging.debug("%s", results)
    return 0, 0


def SpeedTestLatency(name, url):
  """Check latency to a specific SpeedTest.net server.

  Args:
    name: The name of the SpeedTest.net server host.
    url: The URL of the SpeedTest.net server.

  Returns:
    An integer that represents the latency in ms to the SpeedTest.net server.
  """
  logging.info("Checking latency to %s using SpeedTest.net.", name)
  average_time, count, total = 0, 10, 0
  test_url = url+"latency.txt?x=" + str(time.time())
  for _ in range(count):
    error = 0
    start = int(round(time.time() * 1000))
    try:
      urllib2.urlopen(test_url, None, 15)
    except:
      error = 1
    if error == 0:
      finish = int(round(time.time() * 1000))
      average_time += finish - start
      total += 1
    if total == 0:
      return 0
  latency = average_time / total
  logging.info("SpeedTest latency check to %s was %d ms.", name, latency)
  return latency


def SpeedTestBandwidth(name, url):
  """Test the upload and download bandwidth to a specific speedtest.net server.

  DISCLAIMER:  There are some flaws in the methodology used to perform the
  SpeedTest.net tests, however for a local test it is often the best option.
  By local test in the case this script was intended for is testing connectivity
  within the country of Indonesia.  Because Measurement Lab does not have a
  server connected to the local internet exchange, it would be impossible to get
  a clear view of the state of local connectivity as all of the traffic would be
  routed via Singapore, and many providers differentiate local (intra-Indonesia)
  traffic and provide better QoS.

  Args:
    name: Name of the SpeedTest server sponsor.
    url:  URL of the SpeedTest software on sponsor's server.

  Returns:
    Two floats representing the upload and download bandwidth in bits per
    second.
  """
  logging.info("Running SpeedTest bandwidth test to %s.", name)
  st_cmd = "/usr/bin/python tespeed.py -w %s" % url
  results, errors = device.RunCommand(st_cmd, timeout=60, errors=True,
                                      kill_tree=True)
  if errors:
    logging.debug("Encountered an error while running SpeedTest script.")
    return 0, 0
  else:
    speeds = re.search(r"(\d+\.\d+) (\w+) up, (\d+\.\d+) (\w+) down", results)
    if speeds:
      up = speeds.group(1)
      up_unit = speeds.group(2)
      down = speeds.group(3)
      down_unit = speeds.group(4)
      log = ("SpeedTest results to %s: %s %s up and %s %s down." % (name, up,
                                                                    up_unit,
                                                                    down,
                                                                    down_unit))
      logging.info(log)
      if up_unit == "Mbit": up = float(up) * 1048576
      if down_unit == "Mbit": down = float(down) * 1048576
      return int(up), int(down)
    else:
      return 0, 0


def StreamRTSP():
  """Stream a video using RTSP.

  Args:
    None

  Returns:
    1 for pass, 0 for fail and time taken to perform the action
  """
  # TODO(joeyfreeland): I'm using some random video to stream now because
  # YouTube didn't work.  Need to figure out why YouTube didn't work with
  # openRTSP and/or find a more permanent solution for testing.
  logging.info("Trying to stream a video over RTSP.")
  start = time.time()
  rtsp_video = ("rtsp://v3.cache2.c.youtube.com/CjYLENy73wIaLQlEierPwP2jVhMYDSA"
                "NFEIJbXYtZ29vZ2xlSARSBXdhdGNoYLjv7p32_N2wUAw=/0/0/0/video.3gp")
  rtsp_cmd = "/usr/bin/openRTSP -d 30 -D 30 -V %s" % rtsp_video
  errors = device.RunCommand(rtsp_cmd, errors=True, results=False)
  total_time = (time.time() - start) * 1000
  if errors:
    logging.debug("Failed to stream video using RTSP.  Results were:\n %s",
                  errors)
    return 0
  else:
    logging.info("Successfully streamed a video using RTSP. Total time: %d ms.",
                 total_time)
    return total_time
